import os
import subprocess
import paramiko
import logging
from getpass import getpass

# Configure logging
logging.basicConfig(filename='script.log', level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')

def log_message(message):
    """Log a message to both console and log file."""
    print(message)
    logging.info(message)

def run_command(command):
    """Run a shell command and handle errors."""
    try:
        result = subprocess.run(command, check=True, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        output = result.stdout.decode('utf-8')
        log_message(output)
    except subprocess.CalledProcessError as e:
        error_message = f"Error executing command: {e.cmd}\n{e.stderr.decode('utf-8')}"
        log_message(error_message)
        exit(1)

def sync_home_to_remote(remote_user, remote_server, remote_repo_path, password):
    """Sync the local /home directory to the remote repository using rsync over SSH."""
    local_home_path = "/home"
    remote_path = f"{remote_user}@{remote_server}:{remote_repo_path}"
    
    # Use rsync to sync /home to the remote repository
    rsync_command = f"rsync -av --delete {local_home_path}/ {remote_path}/"
    ssh_command = f"sshpass -p {password} {rsync_command}"
    log_message("Syncing /home to remote repository...")
    run_command(ssh_command)

def update_remote_summary(remote_user, remote_server, remote_repo_path, password):
    """Update the summary file on the remote repository using paramiko for SSH."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    log_message("Connecting to remote server to update summary...")
    client.connect(remote_server, username=remote_user, password=password)
    command = f"ostree summary -u --repo={remote_repo_path}"
    stdin, stdout, stderr = client.exec_command(command)
    log_message(stdout.read().decode('utf-8'))
    log_message(stderr.read().decode('utf-8'))
    client.close()

def ensure_mount(mount_point, remote_user, remote_server, remote_repo_path, password):
    """Ensure the remote repository is mounted locally using sshfs and password authentication."""
    if not os.path.ismount(mount_point):
        run_command(f"sudo mkdir -p {mount_point}")
        run_command(f"echo {password} | sudo -S sshfs {remote_user}@{remote_server}:{remote_repo_path} {mount_point}")
        log_message(f"Mounted remote repository {remote_repo_path} at {mount_point}")

def configure_remote(repo_path, remote_name, mount_point):
    """Configure the local repository to include the remote repository."""
    config_file = os.path.join(repo_path, 'config')
    remote_config = f"[remote \"{remote_name}\"]\nurl=file://{mount_point}\ngpg-verify=false\n"
    
    if not os.path.exists(config_file):
        log_message(f"Config file {config_file} does not exist.")
        exit(1)
    
    with open(config_file, 'r') as f:
        config_content = f.read()
    
    if remote_name not in config_content:
        with open(config_file, 'a') as f:
            f.write(remote_config)
        log_message(f"Added remote configuration for {remote_name}")
    else:
        log_message(f"Remote {remote_name} configuration already exists.")

def pull_and_upgrade(repo_path, remote_name, branch_name):
    """Pull updates from the remote repository and perform system upgrade."""
    # Pull the remote branch
    log_message(f"Pulling updates from remote branch {branch_name}...")
    run_command(f"sudo ostree pull --repo={repo_path} {remote_name} {branch_name}")
    
    # Perform system upgrade
    log_message("Upgrading the system...")
    run_command("sudo ostree admin upgrade")

def rollback_system():
    """Rollback the system to the previous deployment."""
    # Perform system rollback
    log_message("Rolling back the system to the previous deployment...")
    run_command("sudo ostree admin rollback")

def main():
    remote_user = "root"
    remote_server = "101.201.181.98"
    remote_repo_path = "/srv/ostree/repo"
    local_repo_path = "/ostree/repo"
    remote_name = "origin1"
    mount_point = "/mnt/remote-ostree"
    branch_name = "testbranch"
    
    # Get password from the user
    password = getpass("Enter the password for the remote server: ")
    
    # Sync /home to the remote repository
    sync_home_to_remote(remote_user, remote_server, remote_repo_path, password)
    
    # Update the summary file on the remote repository
    update_remote_summary(remote_user, remote_server, remote_repo_path, password)

    # Ensure remote repository is mounted
    ensure_mount(mount_point, remote_user, remote_server, remote_repo_path, password)
    
    # Initialize local repository if not already done
    if not os.path.exists(local_repo_path):
        run_command(f"sudo mkdir -p {local_repo_path}")
        run_command(f"sudo ostree init --repo={local_repo_path} --mode=bare")
    
    # Configure remote repository
    configure_remote(local_repo_path, remote_name, mount_point)
    
    # Pull updates and upgrade the system
    pull_and_upgrade(local_repo_path, remote_name, branch_name)

    # Optionally, rollback the system if needed
    rollback_system()

if __name__ == "__main__":
    main()

